function savings_EWM_quadrant = generate_results_capshare_quadrant_rule(bsimax,sample,share,wave,tcost,covariate,winter_prevuse,SMC,savedir,fbase)
% This function calculates cost or energy savings from EWM quadrant 
% rules under a capped budget using grid search. 

% Inputs: 
% (1) bsimax: number of bootstrap runs
% (2) sample: cross-sectional dataset with electricity consumption, 
% covariates, and propensity scores
% (3) share: share of households allowed to treat under the budget cap
% (4) wave: indicates whether the output is wave-specific (=3,6,or 7) or
% for the pooled sample (=0)
% (5) tcost: private marginal cost of implementing the program per household
%  >0 represents cost savings; =0 represents kWh reduction 
% (6) covariate: indicates covariate for analysis: income, size, vintage,
% minimum of baseline consumption, maximum of baseline consumption, or
% standard deviation of consumption
% (7) winter_prevuse: indicates whether baseline consumption is calculated 
% as the mean of consumption in winter months (Jan and Feb) or as the mean
% of specified pre-treatment periods.
% (8) SMC: =1 use social marginal cost; =0 use retail electricity price
% (9) savedir: directory to save the savings table 
% (10) fbase: string used to indicate the propensity score and baseline
% months specification for output table 

% Output: 
% (1) savings_EWM_quadrant: a table summarizing the savings point estimate,
% share of households to treat

%% Gather inputs
[X,Y,D,n,ps]=generate_inputs_quadrant_rule(sample,wave,tcost,covariate,SMC);

%% Estimate g using IPW
g = Y.*(D./ps - (1-D)./(1-ps)); % elements of W(G) - W(G_0)

%% Compress data by histogram

% separation in the two dimensions
if covariate=="income"
    dx1 = 5;    % step size is 5000 for income
    dx2 = 10;   % step size is 10 for pre-treatment consumption
elseif covariate=="size"
    dx1 = 100;
    dx2 = 10;
elseif covariate=="vintage"
    dx1=10;
    dx2=10;
elseif covariate=="min"
    dx1 = 10; % step size is 10 for min/max/std of baseline consumption
    dx2 = 10; % step size is 10 for mean of baseline consumption
elseif covariate=="max"
    dx1 = 10; 
    dx2 = 10; 
elseif covariate=="std"
    dx1 = 10; 
    dx2 = 10; 
end

% sort rows of X -> [Xr, gr]
[Xr, Ind] = sortrows(X); % sorts rows by column (first to last col)
gr = g(Ind);

% use histcounts2 to assign each household to its bin in covariate space
[~,boo1,boo2,xx1,xx2] = histcounts2(Xr(:,1),Xr(:,2),...
    'BinWidth',[dx1,dx2],'Normalization','count');
%boo1: the lower bound of each bin for covariate
%boo2: the lower bound of each bin for prev_use
%xx1: each data point belongs to which covariate bin 
%xx2: each data point belongs to which prev_use bin 

% aggregate data to the bin level
nu = length(boo1)*length(boo2); % # grid points = # cov bins * # prev_use bins
Xu = nan(nu,2); gu = nan(nu,1); nw = nan(nu,1);
k = 1;
for i = 1:length(boo1)
   for j = 1:length(boo2)
       bin_boo = (xx1==i & xx2==j);  % identify households in bin ij
       Xu(k,:) = [boo1(i), boo2(j)]; % coarsened X for bin ij (lower bound)
       gu(k,1) = sum(gr(bin_boo));   % sum gr w/in bin ij
       nw(k,1) = sum(bin_boo);       % num households in bin ij
       k=k+1;
   end
end

% create x1, x2 variables for convenience
x1 = Xu(:,1);
x2 = Xu(:,2);

%% Minimize cost
tic1 = tic; % time grid search

% initialize variables for grid search
grid_x1 = boo1'; grid_x2 = boo2';
k1 = size(grid_x1,1); k2 = size(grid_x2,1);
m_W = Inf*ones(k1,k2);
share_candidate=nan(k1,k2);
fractn=nan(k1,k2);

minW = Inf;

% search grid for optimal quadrant rule
for i1=1:k1
    boo_i1=(x1 - grid_x1(i1));
    for i2=1:k2
        boo_i2=(x2 - grid_x2(i2));
        for sign1 = -1:2:1 %for i = [-1,1]
            for sign2 = -1:2:1
                % treat if both threshold conditions are met
                select = ((sign1.*boo_i1)>=0).*((sign2.*boo_i2)>=0); 
                treat_percent_candidate = sum(nw.*select)/n; % share to treat of the candidate rule 
                fraction = min(1,share./treat_percent_candidate); % rationing term 
                W_rationed = sum(gu.*select).*fraction; % total savings under rationing 
                
                if W_rationed<m_W(i1,i2) % min cost: find lowest W quadrant at i1,i2
                   m_W(i1,i2) = W_rationed; 
                   share_candidate (i1,i2) = treat_percent_candidate;
                   fractn (i1,i2) = fraction;
                end
                
                if (W_rationed < minW) % min cost: find lowest W quadrant overall
                    min_i1 = i1;
                    min_i2 = i2;
                    min_sign1 = sign1;
                    min_sign2 = sign2;
                    minW = W_rationed;        
                end
            end
        end
    end
end

fprintf('Grid search takes %.2f sec\n',toc(tic1))

%% Create variables for savings output (at bottom after bootstrap)

% treatment status for each unique household 
in_Ghat = (min_sign1*(x1 - grid_x1(min_i1))>=0).*(min_sign2*(x2 - grid_x2(min_i2))>=0);
select_fraction=fractn(min_i1,min_i2);
percent=sum(nw.*in_Ghat)/n*select_fraction; % share to treat under rationing
savings=sum(gu.*in_Ghat)/n*select_fraction; % point estimate under rationing

% get in_Ghat from unique household level back to household level 
    in_Ghat_hh = nan(size(g)); % Treatment Assignment from Optimization
    i_u=1;

    for i=1:length(boo1)
        for j=1:length(boo2)
            % individuals with the same (x1,x2)
            bin_boo = logical((xx1==i).*(xx2==j));
            % optimal treatment for people with (x1,x2)
            G_ = in_Ghat(i_u);
            % let all individuals have this assignment
            in_Ghat_hh(bin_boo,1) = G_; % in_Ghat_hh is sorted=> corr to (Xr, gr)
            % step to the next element of Xu/gu/in_Ghat
            i_u=i_u+1;
        end
    end
    
% grid indices
v_x1 = grid_x1;
v_x2 = grid_x2;

if bsimax==0 % if no bootstrap, initiliaze vhats so table code doesn't fail
    vhats_n=0;
    vhats_p=0;
end

 %% Export rule parameters for graphs 

switch winter_prevuse 
    case 1
        s_winter = '_winter';
    case 0
        s_winter = '';
end
if (tcost)>0 % cost_savings
    if (SMC)
        s_cost = 'SMC';
    else
        s_cost = 'PMC';
    end
elseif (tcost)==0
    s_cost='kwh';
end
filename_coefs=sprintf('coef_capshare_quadrant_%s_%s%s_wave%1.0f.mat',covariate,s_cost,s_winter,wave);
filename = sprintf('%s_%s',fbase,filename_coefs);
fpathf = fullfile(savedir,filename);
fprintf('Saving "%s" to "%s" ... ',filename,savedir); 
save(fpathf,'min_i1','min_i2','min_sign1','min_sign2','m_W',...
   'in_Ghat','v_x1','v_x2','Xu','nw','gu','n','covariate','vhats_n','vhats_p','minW',...
   'percent','savings','share_candidate','fractn','select_fraction');
fprintf('   Done!\n');

%% Output for table 

savings_EWM_quadrant=nan(1,4);
savings_EWM_quadrant=array2table(savings_EWM_quadrant,'VariableNames',{'rules','Covariates','Share\ treated','Net\ cost\ changes'}); 
format short g
savings_EWM_quadrant.(1) = {'quadrant'};
savings_EWM_quadrant.(2) = {covariate};
savings_EWM_quadrant.(3) =round(percent*100,2);
savings_EWM_quadrant.(4) =round(savings,3);

end
